#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/epoll.h>
#include <time.h>
#include <stdatomic.h>
#include <stdbool.h>
#include <signal.h>

#include "file_write.h"
#include "data_buffer.h"
#include "log_service.h"

static int sockfd;
static atomic_bool terminate = false;
static pthread_mutex_t addr_lock;
static struct sockaddr_in trans_addr;

static int command_check(const char *data, size_t len, const struct sockaddr_in *addr)
{
    if (!data || !addr) {
        return -1;
    }

    const uint8_t set_trans_addr_cmd[] = {3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8};
    if (len == sizeof(set_trans_addr_cmd) && !memcmp(set_trans_addr_cmd, data, len)) {
        pthread_mutex_lock(&addr_lock);
        memcpy(&trans_addr, addr, sizeof(trans_addr));
        pthread_mutex_unlock(&addr_lock);
        return 0;
    }

    const uint8_t clr_trans_add_cmd[] = {2, 9, 9, 7, 9, 2, 4, 5, 8};
    if (len == sizeof(clr_trans_add_cmd) && !memcmp(clr_trans_add_cmd, data, len)) {
        pthread_mutex_lock(&addr_lock);
        memset(&trans_addr, 0, sizeof(trans_addr));
        pthread_mutex_unlock(&addr_lock);
        return 0;
    }

    return -1;
}

static void *receiver_thread(void *arg)
{
    struct sockaddr_in client_addr;
    socklen_t addr_len = sizeof(client_addr);

    int epoll_fd = epoll_create1(0);
    if (epoll_fd < 0) {
        perror("epoll_create1");
        return NULL;
    }

    struct epoll_event ev;
    ev.events = EPOLLIN;
    ev.data.fd = sockfd;
    if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, sockfd, &ev) < 0) {
        perror("epoll_ctl");
        close(epoll_fd);
        return NULL;
    }

    struct epoll_event events[EPOLL_MAX_EVENTS];
    static char temp_buf[MAX_MSG_SIZE];

    while (!atomic_load(&terminate)) {
        int nfds = epoll_wait(epoll_fd, events, EPOLL_MAX_EVENTS, 100);
        if (nfds < 0) {
            if (errno == EINTR)
                continue;
            perror("epoll_wait");
            break;
        }

        for (int i = 0; i < nfds; i++) {
            if (events[i].data.fd == sockfd) {
                while (!atomic_load(&terminate)) {
                    ssize_t n = recvfrom(sockfd,
                                         temp_buf,
                                         MAX_MSG_SIZE,
                                         /* MSG_TRUNC | */ MSG_DONTWAIT,
                                         (struct sockaddr *)&client_addr,
                                         &addr_len);
                    if (n <= 0) {
                        break;
                    }
                    if (command_check(temp_buf, n, &client_addr)) {
                        data_buffer_t *buf = get_buffer();
                        if (buf) {
                            memcpy(buf->data, temp_buf, n);
                            buf->length = n;
                            enqueue_message(buf);
                        } else {
                            /* drop message, do nothing */
                        }
                    } else {
                        /* this is a command, do not save or re-trans */
                    }
                }
            }
        }
    }

    close(epoll_fd);
    return NULL;
}

static int udp_write(const void *buf, size_t count)
{
    struct sockaddr_in remote;
    pthread_mutex_lock(&addr_lock);
    memcpy(&remote, &trans_addr, sizeof(remote));
    pthread_mutex_unlock(&addr_lock);

    /* if udp re-transportation is not activated, do nothing */
    if (remote.sin_port == 0) {
        return 0;
    }

    const char *p = buf;
    size_t remaining = count;
    while (remaining > 0) {
        ssize_t n = sendto(sockfd, p, remaining, 0, (struct sockaddr *)&remote, sizeof(remote));
        if (n < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
                continue;
            return -2;
        }
        p += n;
        remaining -= n;
    }

    return 0;
}

static void *writer_thread(void *arg)
{
    int result_udp = 0;
    int result_file = file_write_init();
    if (result_file != 0) {
        fprintf(stderr, "file_write_init error: %d.\n", result_file);
        return NULL;
    }
    while (!atomic_load(&terminate)) {
        data_buffer_t *msg = dequeue_message_timed(300);
        if (msg) {
            result_file = file_write(msg->data, msg->length);
            result_udp = udp_write(msg->data, msg->length);
            return_buffer(msg);
            if (result_file != 0) {
                file_write_deinit();
                fprintf(stderr, "file_write error: %d.\n", result_file);
                return NULL;
            }
            if (result_udp != 0) {
                fprintf(stderr, "udp_write error: %d.\n", result_udp);
                // return NULL;
            }
        }
        result_file = file_write_check();
        if (result_file != 0) {
            file_write_deinit();
            fprintf(stderr, "file_write_check error: %d.\n", result_file);
            return NULL;
        }
    }
    file_write_flush();
    file_write_deinit();
    return NULL;
}

static void sigint_handler(int sig)
{
    atomic_store(&terminate, true);
    notify_queue_message();
}

static void cleanup(void)
{
    deinit_buffer_pool();
    deinit_message_queue();
    if (sockfd >= 0) {
        shutdown(sockfd, SHUT_RDWR);
        close(sockfd);
    }
    pthread_mutex_destroy(&addr_lock);
}

int log_service_main(int argc, char *argv[])
{
    memset(&trans_addr, 0, sizeof(trans_addr));
    if (pthread_mutex_init(&addr_lock, NULL) != 0) {
        fprintf(stderr, "trans_addr mutex_lock init error\n");
        return 1;
    }

    int port = LOG_SERVER_PORT;
    struct sockaddr_in server_addr;

    sockfd = socket(AF_INET, SOCK_DGRAM, 0);
    if (sockfd < 0) {
        perror("socket");
        return 1;
    }
    int opt = 1;
    if (setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt))) {
        perror("Setsockopt SO_REUSEADDR failed");
        return 1;
    }

    int flags = fcntl(sockfd, F_GETFL, 0);
    if (fcntl(sockfd, F_SETFL, flags | O_NONBLOCK) < 0) {
        perror("fcntl");
        close(sockfd);
        return 1;
    }

    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(port);

    if (bind(sockfd, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("bind");
        close(sockfd);
        return 1;
    }

    init_buffer_pool();
    init_message_queue();

    pthread_t write_tid;
    if (pthread_create(&write_tid, NULL, writer_thread, NULL) != 0) {
        perror("pthread_create");
        cleanup();
        return 1;
    }

    printf("Server running on port %d. Press Ctrl+C to exit...\n", port);
    signal(SIGINT, sigint_handler);

    receiver_thread(NULL);
    pthread_join(write_tid, NULL);

    cleanup();
    return 0;
}
